{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "from collections import OrderedDict\n",
    "import heapq\n",
    "from typing import (\n",
    "    Hashable,\n",
    "    Callable,\n",
    "    Dict,\n",
    "    Iterable,\n",
    "    Iterator,\n",
    "    List,\n",
    "    Mapping,\n",
    "    MutableMapping,\n",
    "    Optional,\n",
    "    Set,\n",
    "    Tuple,\n",
    "    Union\n",
    ")\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class _Node:\n",
    "    \n",
    "    def __init__(self, key: Hashable, point: np.ndarray) -> None:\n",
    "        self.key = key\n",
    "        self.point = point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vamana(MutableMapping):\n",
    "\n",
    "    def __init__(\n",
    "        self, \n",
    "        dataset: np.ndarray, \n",
    "        alpha: float, \n",
    "        search_list_size: int, \n",
    "        max_degree: int,\n",
    "        distance_func: Callable[[np.ndarray, np.ndarray], float]\n",
    "        ) -> None:\n",
    "        self._dataset = dataset\n",
    "        self._alpha = alpha\n",
    "        self._search_list_size = search_list_size\n",
    "        self._max_degree = max_degree\n",
    "        self._distance_func = distance_func\n",
    "        self._graph: Dict[Hashable, Dict[Hashable, float]] = {}\n",
    "        self._nodes: Dict[Hashable, _Node] = {}\n",
    "\n",
    "\n",
    "    \n",
    "    def greedy_search(\n",
    "            self, \n",
    "            start_node: Hashable, \n",
    "            query: np.ndarray, \n",
    "            topk: int, \n",
    "            search_list_size: int\n",
    "        ) -> Tuple[List[Tuple[float, Hashable]], Set[Hashable]]:\n",
    "        if start_node not in self._nodes:\n",
    "            raise KeyError(start_node)\n",
    "        dist = self._distance_func(self._nodes[start_node].point, query)\n",
    "        search_list = [(dist, start_node)] # 维护search_size\n",
    "        diff_list = [(dist, start_node)] # 维护没有search list - visited list\n",
    "        visited = set()\n",
    "        while diff_list:\n",
    "            dist, p = heapq.heappop(diff_list)\n",
    "            visited.add(p)\n",
    "            dists = [(self._distance_func(self._nodes[out].point, query), out) \n",
    "                     for _, out \n",
    "                     in self._graph[p]\n",
    "                     if out not in visited]\n",
    "            #! 这里应该存在一个更好的算法, 以便于保存search_size个最小\n",
    "            search_list += dists\n",
    "            \n",
    "            if len(search_list) > search_list_size:\n",
    "                search_list = heapq.nsmallest(search_list, search_list_size)\n",
    "            \n",
    "            diff_list = [(d, p) for d, p in search_list if p not in visited]\n",
    "        \n",
    "        return heapq.nsmallest(topk, search_list), visited\n",
    "    \n",
    "    def robust_prune(\n",
    "            self, \n",
    "            key: Hashable, \n",
    "            v: Set[Hashable], \n",
    "            alpha: float, \n",
    "            max_degree: int\n",
    "        ) -> None:\n",
    "        out = [p for _, p in self._graph[key]]\n",
    "        v.update(out)\n",
    "        v.remove(key)\n",
    "        candidates = [\n",
    "            (\n",
    "                self._distance_func(self._nodes[key].point, self._nodes[c].point),\n",
    "                c\n",
    "            )\n",
    "            for c in v\n",
    "        ]\n",
    "        heapq.heapify(candidates)\n",
    "        self._graph[key] = {}\n",
    "\n",
    "        while candidates:\n",
    "            d, p = candidates[0]\n",
    "            self._graph[key][p] = d\n",
    "            if len(self._graph[key]) == max_degree:\n",
    "                break\n",
    "            \n",
    "            candidates = list(filter(lambda dhat, phat: alpha * self._distance_func(\n",
    "                self._nodes[phat].point, self._nodes[p].point) > dhat, candidates))\n",
    "        \n",
    "    \n",
    "    def vamana_index(self):\n",
    "        # 初始化一个R度的有向图\n",
    "        import random\n",
    "        datasize = len(self._dataset)\n",
    "        s = random.randint(0, datasize - 1)\n",
    "        sigma = list(range(datasize))\n",
    "        random.shuffle(sigma)\n",
    "        for i in range(datasize):\n",
    "            self._graph[i] = {\n",
    "                j: self._distance_func(self.data[i], self.data[j]) \n",
    "                for j in random.sample(range(datasize), self._max_degree)}\n",
    "        for i in sigma:\n",
    "            L, V = self.greedy_search(s, self._dataset[i], 1, self._search_list_size)\n",
    "            self.robust_prune(i, V, self._alpha, self._max_degree)\n",
    "            for j, d in self._graph[i]:\n",
    "                if len(self._graph[j]) + 1 > self._max_degree:\n",
    "                    self.robust_prune(j, [v for _, v in self._graph[j] if v != i] + [i], self._alpha, self._max_degree)\n",
    "                else:\n",
    "                    self._graph[j][i] = d        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ann",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
